import os
import asyncio
import re

from edge_tts.exceptions import NoAudioReceived
from util.file_util import merge_audio_file

from util.tts_utils import tts


def split_text_into_batches(text, max_tokens_per_batch):
    # 定义一个正则表达式，在中文标点符号处拆分句子
    sentence_splitter = re.compile(r'(?<=[。！？])')

    # 将文本拆分为句子
    sentences = [sentence.strip() for sentence in sentence_splitter.split(text) if sentence.strip()]

    # 初始化变量
    batches = []
    current_batch = ""

    for sentence in sentences:
        if len(current_batch) + len(sentence) <= max_tokens_per_batch:
            current_batch += sentence + " "
        else:
            # 找到距离 max_tokens_per_batch 限制最近的标点符号
            last_punctuation_index = max(current_batch.rfind('。'), current_batch.rfind('！'), current_batch.rfind('？'))

            # 如果限制范围内没有标点符号，就在最后一个空格处拆分
            split_index = last_punctuation_index if last_punctuation_index != -1 else current_batch.rfind(' ')

            # 将批次添加到拆分索引处
            batches.append(current_batch[:split_index].strip())

            # 新批次从拆分索引开始
            current_batch = sentence + " "

    if current_batch.strip():  # 确保不将空字符串添加到批次中
        batches.append(current_batch.strip())

    return batches


async def traverse_dir(path):
    for root, dirs, files in os.walk(path):
        print("当前目录：", root)
        print("子目录列表：", dirs)
        print("文件列表：", files)

        for file in files:
            i = 1
            with open(path + '/' + file, encoding='utf-8') as f:
                dest_path = (f'C:\\Users\\70745\\PycharmProjects\\'
                             f'edge-tts-wrapper\\tts/{file}.mp3')
                mp3_list = []
                merge_flg = False
                for content_line in f:
                    lines = split_text_into_batches(content_line, 800)
                    for line in lines:
                        mp3_path = (f'C:\\Users\\70745\\PycharmProjects\\'
                                    f'edge-tts-wrapper\\tts/{file}{"{:0>2d}".format(i)}.mp3')
                        try:
                            await tts(line, mp3_path)
                            i = i + 1
                            mp3_list.append(mp3_path)
                            if len(mp3_list) > 10:
                                merge_audio_file(mp3_list, dest_path)
                                merge_flg = True
                                mp3_list.clear()
                                os.rename(dest_path, f'{dest_path}tmp.mp3')
                                mp3_list.append(f'{dest_path}tmp.mp3')
                        except NoAudioReceived as e:
                            print(e)
                if (merge_flg and len(mp3_list) > 1) or len(mp3_list) > 0:
                    merge_audio_file(mp3_list, dest_path)


if __name__ == "__main__":
    loop = asyncio.get_event_loop_policy().get_event_loop()
    try:
        dir_path = "C:\\Users\\70745\\PycharmProjects\\edge-tts-wrapper\\data"
        print('待遍历的目录为：', dir_path)
        print('遍历结果为：')
        loop.run_until_complete(traverse_dir(dir_path))
    finally:
        loop.close()
